#include<assert.h>
#include<iostream>

namespace xzj
{
    enum Color
    {
        RED,
        BLACK
    };

    template<class K>
    struct RBTreeNode 
    {
        RBTreeNode(const K& key)
            :_parent(nullptr)
            ,_left(nullptr)
            ,_right(nullptr)
            ,_key(key)
            ,_col(RED)
        {}

        RBTreeNode* _parent;
        RBTreeNode* _left;
        RBTreeNode* _right;
        K _key;
        Color _col;
    };

    template<class K>
    class RBTree
    {
        using Node = RBTreeNode<K>;
    public:
        bool insert(const K& key)
        {
            if(_root == nullptr)
            {
                _root = new Node(key);
                _root->_col = BLACK;
                return true;
            }

            Node* cur = _root;
            Node* parent = nullptr;
            while(cur)
            {
                parent = cur;
                if(cur->_key>key)
                {
                    cur = cur->_left;
                }
                else if(cur->_key<key)
                {
                    cur = cur->_right;
                }
                else 
                {
                    return false;
                }
            }

            cur = new Node(key);
            if(parent->_key>key)
            {
                parent->_left = cur;
            }
            else 
            {
                parent->_right = cur;
            }
            cur->_parent = parent;
            
            while(parent&&parent->_col==RED)
            {
                Node* grandparent = parent->_parent;
                if(grandparent->_left==parent)
                {
                    Node* uncle = grandparent->_right;
                    if(uncle&&uncle->_col==RED)
                    {
                        parent->_col = BLACK;
                        uncle->_col = BLACK;
                        grandparent->_col = RED;
                        parent = grandparent->_parent;
                        cur = grandparent;
                    }
                    else 
                    {

                         if(cur==parent->_left)
                         {
                             RotateR(grandparent);
                             parent->_col = BLACK;
                             grandparent->_col = RED;
                         }
                         else if(cur==parent->_right)
                         {
                             RotateL(parent);
                             RotateR(grandparent);

                             cur->_col = BLACK;
                             grandparent->_col = RED;
                         }
                         else 
                         {
                             assert(false);
                         }

                         break;
                    }
                }
                else 
                {
                    Node* uncle = grandparent->_left;
                    
                    if(uncle&&uncle->_col==RED)
                    {
                        parent->_col = BLACK;
                        uncle->_col = BLACK;
                        grandparent->_col = RED;
                        parent = grandparent->_parent;
                        cur = grandparent;
                    }
                    else 
                    {

                         if(cur==parent->_right)
                         {
                             RotateL(grandparent);
                             parent->_col = BLACK;
                             grandparent->_col = RED;
                         }
                         else if(cur==parent->_left)
                         {
                             RotateR(parent);
                             RotateL(grandparent);

                             cur->_col = BLACK;
                             grandparent->_col = RED;
                         }
                         else 
                         {
                             assert(false);
                         }

                         break;
                    }
                }
                
            }

            _root->_col = BLACK;
                
            return true;
        }

        bool isRBTreeBalance()
        {
            //检查根
            if(_root&&_root->_col==RED) return false;

            Node* cur = _root;
            int pathNum = 0;//路径黑色结点数目
            while(cur)
            {
                if(cur->_col==BLACK) pathNum++;

                cur = cur->_left;
            }

            int BNodecount = 0;//计算到每条路径末尾的黑色结点数目
            return _isRBTreeBalance(_root,pathNum,BNodecount);
        }

        void InOrder()
        {
            _InOrder(_root);
        }

        void preOrder()
        {
            _preOrder(_root);
        }

    private:
        void RotateL(Node* parent)
        {
            Node* subR = parent->_right;
            Node* subRL = subR->_left;

            if(subRL)
                subRL->_parent = parent;
            parent->_right = subRL;
            subR->_left = parent;

            if(parent->_parent==nullptr)
            {
                _root = subR;
                subR->_parent=nullptr;
            }
            else 
            {
                if(parent->_parent->_left==parent)
                {
                    parent->_parent->_left = subR;
                }
                else 
                {
                    parent->_parent->_right = subR;
                }
                subR->_parent = parent->_parent;
            }
            parent->_parent = subR;
        }

        void RotateR(Node* parent)
        {
            Node* subL = parent->_left;
            Node* subLR = subL->_right;

            if(subLR)
                subLR->_parent = parent;
            parent->_left = subLR;
            subL->_right = parent;

            if(parent->_parent==nullptr)
            {
                _root = subL;
                subL->_parent = nullptr;
            }
            else 
            {
                if(parent->_parent->_left==parent)
                {
                    parent->_parent->_left = subL;
                }
                else 
                {
                    parent->_parent->_right = subL;
                }
                subL->_parent = parent->_parent;
            }
            parent->_parent = subL;
        }

        bool _isRBTreeBalance(Node* root, int pathNum,int BNodecount)
        {
            if(root==nullptr)
            {
                if(BNodecount==pathNum) return true;
                else return false;
            }

            if(root->_col==RED&&root->_parent&&root->_parent->_col==RED)
                return false;

            if(root->_col==BLACK)
                BNodecount++;

            return _isRBTreeBalance(root->_left,pathNum,BNodecount)
                &&_isRBTreeBalance(root->_right,pathNum,BNodecount);
        }

        void _InOrder(Node* root)
        {
            if(root==nullptr) return;

            _InOrder(root->_left);
            std::cout<<root->_key<<" ";
            _InOrder(root->_right);
        }

        void _preOrder(Node* root)
        {
            if(root==nullptr) return;

            std::cout<<root->_key<<":"<<root->_col<<std::endl;
            _preOrder(root->_left);
            _preOrder(root->_right);
        }

        Node* _root=nullptr;
    };
}
